import base64
import mimetypes
from pathlib import Path
import json
import aiohttp
from pydantic import BaseModel, Field
from typing import Optional, List

class Location(BaseModel):
    """商家位置信息"""
    city: Optional[str] = None
    state: Optional[str] = None
    zipcode: Optional[str] = None

class LineItem(BaseModel):
    """购物项目明细"""
    description: Optional[str] = None  # 商品描述
    product_code: Optional[str] = None  # 商品代码/SKU
    category: Optional[str] = None  # 商品类别
    item_price: Optional[str] = None  # 单品价格
    sale_price: Optional[str] = None  # 折扣价格
    quantity: Optional[str] = None  # 数量
    total: Optional[str] = None  # 单项总价

class ReceiptDetails(BaseModel):
    """收据详细信息"""
    merchant: Optional[str] = None  # 商家名称
    location: Optional[Location] = None  # 商家位置，允许为None
    time: Optional[str] = None  # 交易时间
    items: List[LineItem] = Field(default_factory=list)  # 商品列表，默认空列表
    subtotal: Optional[str] = None  # 小计金额
    tax: Optional[str] = None  # 税费
    total: Optional[str] = None  # 总金额
    handwritten_notes: List[str] = Field(default_factory=list)  # 手写备注，默认空列表

# 千问API配置 - 使用兼容模式端点
QWEN_API_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
QWEN_API_KEY = "sk-4e88cf4db3e14894bafaff606d296610"  # 替换为你的千问API密钥

# 更清晰的系统提示词
SYSTEM_PROMPT = """你是一个专业的收据信息提取系统。请仔细分析提供的收据图片，提取所有相关信息并以严格的JSON格式返回。

返回的JSON必须包含以下字段：
{
  "merchant": "商家名称或null",
  "location": {
    "city": "城市或null", 
    "state": "州或省份或null",
    "zipcode": "邮政编码或null"
  },
  "time": "交易时间(ISO格式)或null",
  "items": [
    {
      "description": "商品描述",
      "product_code": "商品代码或null",
      "category": "商品类别或null",
      "item_price": "单品价格",
      "sale_price": "折扣价格或null",
      "quantity": "数量",
      "total": "单项总价"
    }
  ],
  "subtotal": "小计金额或null",
  "tax": "税费或null",
  "total": "总金额或null",
  "handwritten_notes": ["备注1", "备注2"]  // 如果没有手写备注，使用空数组[]
}

重要规则：
1. 如果字段不存在或无法识别，使用null值（除了handwritten_notes必须使用空数组[]）
2. 所有金额格式化为字符串数字（如"12.99"）
3. 日期格式为YYYY-MM-DDTHH:MM:SS
4. 严格返回纯JSON格式，不要包含其他任何文本
5. handwritten_notes字段必须始终是数组，即使为空
"""

async def extract_receipt_details(
        image_path: str,
        model: str = "qwen-vl-max"  # 千问视觉语言模型
) -> ReceiptDetails:
    """
    从收据图片中提取结构化详细信息（使用千问模型兼容模式）
    """

    # 验证文件存在
    if not Path(image_path).exists():
        raise FileNotFoundError(f"图片文件不存在: {image_path}")

    # 确定图片的MIME类型
    mime_type, _ = mimetypes.guess_type(image_path)
    if mime_type is None:
        mime_type = "image/jpeg"

    # 读取图片并进行base64编码
    try:
        image_data = Path(image_path).read_bytes()
        b64_image = base64.b64encode(image_data).decode("utf-8")
    except Exception as e:
        raise IOError(f"读取图片文件失败: {str(e)}")

    # 构建请求头
    headers = {
        "Authorization": f"Bearer {QWEN_API_KEY}",
        "Content-Type": "application/json",
    }

    # 构建请求体
    payload = {
        "model": model,
        "messages": [
            {
                "role": "system",
                "content": SYSTEM_PROMPT
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "请分析这张收据图片并严格按照要求的JSON格式返回提取的信息。"
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:{mime_type};base64,{b64_image}"
                        }
                    }
                ]
            }
        ],
        "max_tokens": 3000,
        "temperature": 0.1,
        "response_format": { "type": "json_object" }
    }

    try:
        # 发送异步请求
        async with aiohttp.ClientSession() as session:
            async with session.post(
                    QWEN_API_URL,
                    headers=headers,
                    json=payload,
                    timeout=60
            ) as response:

                if response.status != 200:
                    error_text = await response.text()
                    raise Exception(f"API请求失败: HTTP {response.status}, 错误信息: {error_text}")

                result = await response.json()

                # 解析响应
                if "choices" in result and len(result["choices"]) > 0:
                    content = result["choices"][0]["message"]["content"]

                    # 清理响应内容
                    content = content.strip()
                    print(f"原始API响应: {content}")  # 调试信息

                    # 处理可能的代码块格式
                    if content.startswith("```json"):
                        content = content.replace("```json", "").replace("```", "").strip()
                    elif content.startswith("```"):
                        content = content.replace("```", "", 1).rsplit("```", 1)[0].strip()

                    # 解析JSON
                    try:
                        json_data = json.loads(content)
                        print(f"解析后的JSON: {json_data}")  # 调试信息

                        # 预处理数据，确保所有字段都有合适的默认值
                        processed_data = preprocess_json_data(json_data)
                        print(f"处理后的数据: {processed_data}")  # 调试信息

                        return ReceiptDetails(**processed_data)
                    except json.JSONDecodeError as e:
                        print(f"JSON解析失败，原始内容: {content}")
                        raise ValueError(f"JSON解析失败: {str(e)}")
                else:
                    raise ValueError("API响应格式不符合预期")

    except aiohttp.ClientError as e:
        raise ConnectionError(f"网络请求失败: {str(e)}")
    except Exception as e:
        raise RuntimeError(f"处理失败: {str(e)}")

def preprocess_json_data(raw_data: dict) -> dict:
    """
    预处理JSON数据，确保所有字段都有合适的格式

    参数:
        raw_data: 原始API返回的JSON数据

    返回:
        处理后的数据字典
    """
    processed = raw_data.copy()

    # 确保handwritten_notes是列表
    if "handwritten_notes" not in processed or processed["handwritten_notes"] is None:
        processed["handwritten_notes"] = []
    elif not isinstance(processed["handwritten_notes"], list):
        # 如果不是列表，尝试转换为列表或使用空列表
        processed["handwritten_notes"] = []

    # 确保items是列表
    if "items" not in processed or processed["items"] is None:
        processed["items"] = []
    elif not isinstance(processed["items"], list):
        processed["items"] = []

    # 确保location是字典或None
    if "location" in processed and processed["location"] is None:
        processed["location"] = {"city": None, "state": None, "zipcode": None}
    elif "location" not in processed:
        processed["location"] = {"city": None, "state": None, "zipcode": None}

    # 清理items中的None值
    if processed["items"]:
        for item in processed["items"]:
            for key in list(item.keys()):
                if item[key] is None:
                    item[key] = None

    return processed

# 测试函数
async def test_extraction():
    """测试收据信息提取"""
    try:
        # 示例收据路径 - 请替换为实际路径
        example_receipt = "D:/tmp/example_receipt.jpg"

        # 检查文件是否存在
        if not Path(example_receipt).exists():
            print("请提供有效的收据图片路径")
            return

        print("开始提取收据信息...")

        # 提取收据信息
        result = await extract_receipt_details(example_receipt)

        # 打印结果
        print("\n=== 收据提取结果 ===")
        print(f"商家: {result.merchant}")
        if result.location:
            print(f"位置: {result.location.city}, {result.location.state} {result.location.zipcode}")
        print(f"时间: {result.time}")
        print(f"小计: {result.subtotal}")
        print(f"税费: {result.tax}")
        print(f"总计: {result.total}")

        print(f"\n商品数量: {len(result.items)}")
        for i, item in enumerate(result.items, 1):
            print(f"\n商品 {i}:")
            print(f"  描述: {item.description}")
            print(f"  类别: {item.category}")
            print(f"  原价: {item.item_price}")
            print(f"  折扣价: {item.sale_price}")
            print(f"  数量: {item.quantity}")
            print(f"  总价: {item.total}")

        print(f"\n手写备注: {result.handwritten_notes}")

    except FileNotFoundError as e:
        print(f"文件错误: {str(e)}")
    except ConnectionError as e:
        print(f"网络连接错误: {str(e)}")
        print("请检查API密钥和网络连接")
    except ValueError as e:
        print(f"数据解析错误: {str(e)}")
    except Exception as e:
        print(f"未知错误: {str(e)}")
        import traceback
        traceback.print_exc()  # 打印完整的错误堆栈

# 如果需要运行测试，取消注释下面的代码
import asyncio
asyncio.run(test_extraction())